/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Contains declaration of the grain_collection_grain_model class. */

// $Id$

#ifndef __grain_collection_grain_model__
#define __grain_collection_grain_model__

#include "grain_model.h"
#include "grain.h"
#include "thermal_equilibrium_grain.h"
#include "grain_hack.h"
#include "boost/shared_ptr.hpp"
#include "preferences.h"
#include "boost/foreach.hpp"

namespace mcrx {
  template <template<class> class, typename> class grain_collection_grain_model;
};


/** This class is a grain model made up from a collection of grain
    objects. In this case, it's easy to calculate the extinction curve
    so we collect these in a base class. The concrete classes deriving
    from this must initialize the size distributions, since those are
    arbitrary. */
template <template<class> class chromatic_policy, 
	  typename T_rng_policy = mcrx::global_random> 
class mcrx::grain_collection_grain_model : 
  public mcrx::grain_model<chromatic_policy, T_rng_policy> {
protected:
  std::vector<boost::shared_ptr<emitting_grain> > grains_;
  std::vector<array_1> dnda_;

  /// Unit conversion for SED calc
  T_float sed_unit_conversion_;

  /** Calculates the extinction curve of the grain model. This must be
      called from the derived classes once the grains_, dnda_ and
      units have been loaded. */
  void calculate_extinction_curve();

public:
  grain_collection_grain_model() : 
    grain_model<chromatic_policy, T_rng_policy>(),
    // Initialize to NaN so we don't use unitialized 
    sed_unit_conversion_(blitz::quiet_NaN(T_float())) {};
  grain_collection_grain_model (const grain_collection_grain_model& g) :
    grain_model<chromatic_policy, T_rng_policy>(g),
    sed_unit_conversion_(g.sed_unit_conversion_) {
    // need to make sure array members are copied. We make the copies
    // thread-local, since the methods are not reentrant anyway.
    BOOST_FOREACH (boost::shared_ptr<emitting_grain> gg, g.grains_) {
      grains_.push_back(gg->clone());
    }
    BOOST_FOREACH (array_1 a, g.dnda_) {
      dnda_.push_back(make_thread_local_copy(a));
    }
  };
  virtual boost::shared_ptr<grain_model<chromatic_policy, T_rng_policy> > 
  clone() const {
    return boost::shared_ptr<grain_model<chromatic_policy, T_rng_policy> >
      (new grain_collection_grain_model(*this));};

  /// Returns the wavelength array used for absorption calculations.
  const array_1& alambda() const {return grains_[0]->alambda();};
  /// Returns the wavelength array used for emission calculations.
  const array_1& elambda() const {return grains_[0]->elambda();};

  /** Resamples all the grain objects to use the specified absorption
      and emission wavelength vectors. Also update the extinction
      curve so the two are consistent. The grains use separate
      absorption and emission wavelength vectors so that it is
      possible to calculate the intensities on a low-resolution grid
      while calculating the emission SED on a high-resolution
      one. Note that this function is NOT necessarily safe to call
      several times with the same wavelength vector, because
      resampling with identical points does not necessarily preserve
      the data. \todo is this still true? */
  void resample (const array_1& lambda, const array_1& elambda=array_1()) {
    BOOST_FOREACH (boost::shared_ptr<emitting_grain> g, grains_) {
      g->resample(lambda,elambda);
    }
    calculate_extinction_curve();

    // This is where we set the unit conversions for the SED
    // calculation, because these need to be set up before calculating
    // SEDs (and it's very inefficient to do it for every cell).  The
    // SED of the grains is in SI, but we need to make sure dnda and
    // delta_size is in consistent units. dnda is in scatterer units
    // while delta_size is in grain units, so we convert dnda to grain
    // units.
    sed_unit_conversion_ = 1./units::convert
      ( mcrx::scatterer<chromatic_policy, T_rng_policy>::units().get("length"), 
	grains_[0]->units().get("length"));
  };

  /** This redefinition of the virtual scatterer::set_wavelength
      function just calls the resample function (which also
      recalculates the extinction curve). */
  virtual void set_wavelength (const array_1& lambda) {
    resample(lambda);
  };

  /** This gets called by the template virtual hack in
      grain_model::calculate_SED. */
  template<typename T>
  void calculate_SED_virtual(const blitz::ETBase<T>& intensity, 
			     const array_1& m_dust, array_2& sed) const;

  /** This function returns the total absorption cross section for the
      dust model given a certain dust mass. Intended for testing. */
  array_1 sigma_tot(T_float m_dust) const {
    using namespace blitz;
    array_1 st(grains_[0]->sigma().extent(secondDim));
    st=0;
    for (int i =  0; i <grains_.size(); ++i) {
      st += sum (grains_[i]->sigma()(tensor::j, tensor::i)*
		 dnda_[i](blitz::tensor::j) *
		 grains_[i]->delta_size()(tensor::j), tensor::j)*m_dust;
    }
    return st;
  };

};



template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
void mcrx::grain_collection_grain_model<chromatic_policy, T_rng_policy>::
calculate_extinction_curve()
{
  using namespace blitz;
  using namespace boost;

  std::cout << "Calculating extinction curve" << std::endl;

  assert(grains_.size()>0);
  assert(grains_.size()==dnda_.size());

  // check that grain units and wavelengths are internally consistent
  BOOST_FOREACH(boost::shared_ptr<emitting_grain> g, grains_) {
    check_consistency(grains_.front()->units(), g->units(), true);
    ASSERT_ALL(grains_.front()->alambda()==g->alambda());
  }

  // we need to convert units. The extinction curve should be in the
  // units given by the scatterer units object, normalized to 1 mass
  // unit in that system. In the calculations below, the grain
  // quantities need to be converted. sigma has units of length^2 and
  // deltasize units of length. The conversion is thus lengthconv^3.
  const T_unit_map& us = 
    mcrx::scatterer<chromatic_policy, T_rng_policy>::units();
  const T_float length_conversion = 
    units::convert(grains_.front()->units().get("length"), 
		   us.get("length"));

  // The units of the extinction curve is area/mass, ie we should
  // divide by the mass conversion. But the mass unit is given by the
  // dnda which is in scatterer units, so there is no conversion here.

  // we don't currently convert the wavelength vector, though we want to check
  assert(grains_.front()->units().get("wavelength")==
	 us.get("wavelength"));

  // calculate extinction curve. this is based on the *emission* wavelength grid
  const int n = grains_.front()->elambda().size();
  array_1 abs(n);
  abs = 0;
  for(int i=0; i<grains_.size(); ++i) {
    abs +=
      sum((grains_[i]->esigma()) (tensor::j, tensor::i) * 
	  (dnda_[i])(tensor::j) * 
	  (grains_[i]->delta_size()) (tensor::j), tensor::j);
  };
  abs *= (length_conversion*length_conversion*length_conversion);

  array_1 sca(n);
  sca = 0;
  for(int i=0; i<grains_.size(); ++i) {
    sca +=
      sum((grains_[i]->sigma_sca())(tensor::j, tensor::i) * 
	  (dnda_[i])(tensor::j) * 
	  (grains_[i]->delta_size())(tensor::j), tensor::j);
  };
  sca *= (length_conversion*length_conversion*length_conversion);

  array_1 g(n);
  g = 0;
  for(int i=0; i<grains_.size(); ++i) {
    g +=
      sum((grains_[i]->sigma_sca())(tensor::j, tensor::i) * 
	  (grains_[i]->g())(tensor::j, tensor::i) * 
	  (dnda_[i])(tensor::j) * 
	  (grains_[i]->delta_size())(tensor::j), tensor::j);
      
  };
  g *= where(sca>0,
	     (length_conversion*length_conversion*length_conversion)/sca,
	     0);
  ASSERT_ALL(abs>=0);
  ASSERT_ALL(sca>=0);
  ASSERT_ALL(g>=-1);
  ASSERT_ALL(g<= 1);

  // CHECK THAT THESE VALUES ARE OK

  // now copy these to the scatterer members.
  assign(this->opacity_, (abs+sca));
  assign(this->albedo_, sca/(abs+sca));
  this->phfunc.set_g (g);

  DEBUG(1,								\
	ofstream crap("extinction_curve.txt");				\
	crap << "# units " << us.get("length") << '\t' << us.get("mass") << '\n'; \
	for(int i=0;i<grains_.front()->alambda().size();++i)		\
	  crap << grains_.front()->alambda()(i) << '\t' << (abs(i)+sca(i)) << '\t' \
	       << sca (i)/(abs (i) +sca (i)) << '\t'			\
	       << g (i) << endl;					\
	);
}


template <template<class> class chromatic_policy, 
	  typename T_rng_policy> 
template <typename T>
void
mcrx::grain_collection_grain_model<chromatic_policy, T_rng_policy>::
calculate_SED_virtual(const blitz::ETBase<T>& intensity, 
		      const array_1& m_dust,
		      array_2& sed) const
{
  using namespace blitz;

  // add emission for all grain species. 
  for (int i =  0; i <grains_.size(); ++i) {
    cout << "\t\t Calculating dust emission SED for grain species " << i << endl;
    // we incorporate the SED unit conversion into the dn array
    const array_1 dn(dnda_[i]*
		     grains_[i]->delta_size()*
		     sed_unit_conversion_);
    grains_[i]->calculate_SED_from_intensity(intensity,
					     m_dust,
					     dn,
					     sed,
					     true);
  }
}

#endif
